import torch
from torch.utils import data
from torch import nn
from torch.optim import lr_scheduler
from dataset import custom_dataset
from model import EAST
from loss import Loss
import os
import time
import numpy as np
from utils import load_json

class Trainer:
	def __init__(self,config):
		self.train_img_path=config['train_img_path']
		self.train_gt_path=config['train_gt_path']
		self.pths_path=config['pths_path']
		self.batch_size=config['batch_size']
		self.lr=config['lr']
		self.num_workers=config['num_workers']
		self.epoch_iter=config['epoch_iter']
		self.interval=config['save_interval']
		self.global_step=0

		if self.tensorboard_enable:
			from torch.utils.tensorboard import SummaryWriter
			self.writer = SummaryWriter(self.tensorboard_dir)

	def train(self):
	# def train(train_img_path, train_gt_path, pths_path, batch_size, lr, num_workers, epoch_iter, interval):
		file_num = len(os.listdir(self.train_img_path))
		trainset = custom_dataset(self.train_img_path, self.train_gt_path)
		train_loader = data.DataLoader(trainset, batch_size=self.batch_size, \
									   shuffle=True, num_workers=self.num_workers, drop_last=True)

		criterion = Loss()
		device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
		model = EAST()
		data_parallel = False
		if torch.cuda.device_count() > 1:
			model = nn.DataParallel(model)
			data_parallel = True
		model.to(device)
		optimizer = torch.optim.Adam(model.parameters(), lr=self.lr)
		scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[self.epoch_iter//2], gamma=0.1)

		for epoch in range(self.epoch_iter):
			model.train()

			epoch_loss = 0
			epoch_time = time.time()
			for i, (img, gt_score, gt_geo, ignored_map) in enumerate(train_loader):
				start_time = time.time()
				img, gt_score, gt_geo, ignored_map = img.to(device), gt_score.to(device), gt_geo.to(device), ignored_map.to(device)
				pred_score, pred_geo = model(img)
				loss = criterion(gt_score, pred_score, gt_geo, pred_geo, ignored_map)

				epoch_loss += loss.item()
				optimizer.zero_grad()
				loss.backward()
				optimizer.step()

				if self.tensorboard_enable:
					# write tensorboard
					self.writer.add_scalar('TRAIN/loss_all', loss.item(), self.global_step)


				print('Epoch is [{}/{}], mini-batch is [{}/{}], time consumption is {:.8f}, batch_loss is {:.8f}'.format( \
					self.global_step, self.epoch_iter, i+1, int(file_num/self.batch_size), time.time()-start_time, loss.item()))
				self.global_step+=1

			print('epoch_loss is {:.8f}, epoch_time is {:.8f}'.format(epoch_loss/int(file_num/self.batch_size), time.time()-epoch_time))
			print(time.asctime(time.localtime(time.time())))
			print('='*50)
			if (epoch + 1) % self.interval == 0:
				state_dict = model.module.state_dict() if data_parallel else model.state_dict()
				torch.save(state_dict, os.path.join(self.pths_path, 'model_epoch_{}.pth'.format(epoch+1)))
			scheduler.step()
		if self.tensorboard_enable:
			self.writer.close()



if __name__ == '__main__':


	config = load_json('config.json')
	trainer=Trainer(config)
	trainer.train()

